"""
查询分类器
根据用户查询的特征，将其分类到不同的类别
"""

import re
from typing import Dict, List
from base.logger import logger
from .prompts import QUERY_CLASSIFICATION_TEMPLATE


class QueryClassifier:
    """查询分类器"""
    
    def __init__(self):
        self.classification_rules = {
            'factual': [
                r'什么是', r'谁是', r'哪里', r'何时', r'多少',
                r'定义', r'含义', r'概念'
            ],
            'conceptual': [
                r'为什么', r'如何理解', r'原理', r'机制',
                r'解释', r'说明', r'阐述'
            ],
            'procedural': [
                r'如何', r'怎么', r'步骤', r'方法',
                r'操作', r'实现', r'执行'
            ],
            'comparative': [
                r'比较', r'对比', r'区别', r'差异',
                r'优缺点', r'异同', r'相似'
            ]
        }
    
    def classify_query(self, query: str) -> str:
        """
        对查询进行分类
        
        Args:
            query: 用户查询
            
        Returns:
            查询类别
        """
        query = query.lower().strip()
        
        # 基于规则的分类
        for category, patterns in self.classification_rules.items():
            for pattern in patterns:
                if re.search(pattern, query):
                    logger.info(f"查询分类: {query} -> {category}")
                    return category
        
        # 默认分类
        logger.info(f"查询分类: {query} -> other")
        return 'other'
    
    def get_query_features(self, query: str) -> Dict[str, any]:
        """
        提取查询特征
        
        Args:
            query: 用户查询
            
        Returns:
            查询特征字典
        """
        features = {
            'length': len(query),
            'word_count': len(query.split()),
            'has_question_mark': '?' in query or '？' in query,
            'category': self.classify_query(query),
            'keywords': self._extract_keywords(query)
        }
        
        return features
    
    def _extract_keywords(self, query: str) -> List[str]:
        """
        提取查询中的关键词
        
        Args:
            query: 用户查询
            
        Returns:
            关键词列表
        """
        # 简单的关键词提取（可以后续优化为使用NLP工具）
        stop_words = {'的', '是', '在', '有', '和', '与', '或', '但', '然而', '因为', '所以'}
        words = query.split()
        keywords = [word for word in words if word not in stop_words and len(word) > 1]
        
        return keywords[:5]  # 返回前5个关键词